import torch
import torch.nn as nn

from geomloss import SamplesLoss


def geo_center_pull_loss(list_vec_groups):

    num_group_items = []
    group_loss = []
    for vecs in list_vec_groups:
        num_items = vecs.size(0)
        vec_avg = vecs.sum(0) / num_items
        dist = torch.norm(vecs - vec_avg, dim=1)
        sum_dist = dist.sum()
        num_group_items.append(num_items)
        group_loss.append(sum_dist)
    num_items = sum(num_group_items)
    loss = sum(group_loss)

    return loss / num_items

def geo_center_push_loss(list_vec_groups, sigma):

    num_groups = len(list_vec_groups)
    if num_groups < 2:
        return list_vec_groups[0].sum() * 0
    group_centers = []
    for vecs in list_vec_groups:
        num_items = vecs.size(0)
        vec_avg = vecs.sum(0) / num_items
        group_centers.append(vec_avg)
    group_centers = torch.stack(group_centers, dim=0)
    vec_dims = group_centers.size(1)
    group_centersA = group_centers.view(num_groups, 1, vec_dims).repeat(1, num_groups, 1)
    group_centersB = group_centers.view(1, num_groups, vec_dims).repeat(num_groups, 1, 1)
    center_dist = torch.norm(group_centersA - group_centersB, dim=-1)
    loss = torch.exp(-center_dist / (2 * sigma * sigma)).sum() / (num_groups * num_groups)

    return loss

class FeaturePullLoss(nn.Module):

    def __init__(self, loss_weight=1.0, loss_type='geo'):
        super().__init__()
        self.loss_weight = loss_weight
        self.loss = None
        if loss_type == 'geo':
            self.loss = geo_center_pull_loss

        assert self.loss != None

    def forward(self,
                list_vec_groups):

        loss_pull = self.loss_weight * self.loss(list_vec_groups)

        return loss_pull

class FeaturePushLoss(nn.Module):

    def __init__(self, loss_weight=1.0, sigma=1.0, loss_type='geo'):
        super().__init__()
        self.loss_weight = loss_weight
        self.sigma = sigma
        self.loss = None
        if loss_type == 'geo':
            self.loss = geo_center_push_loss

        assert self.loss != None

    def forward(self,
                list_vec_groups):

        push_loss = self.loss_weight * self.loss(list_vec_groups, self.sigma)

        return push_loss


class FeaturePushLossWasserstein(nn.Module):

    def __init__(self,
                 loss_weight=1.0, 
                 loss_type='sinkhorn',
                 p=1,
                 blur=0.05):
        super().__init__()
        self.loss_weight = loss_weight
        self.loss = SamplesLoss(loss=loss_type, p=p, blur=blur)

    def group_loss(self,
                    list_vec_groups):

        num_groups = len(list_vec_groups)
        if num_groups < 2:
            return list_vec_groups[0].sum() * 0
        group_losses = []
        for i in range(num_groups):
            for j in range(1, num_groups):
                w_loss = self.loss(list_vec_groups[i], list_vec_groups[j])
                group_losses.append(w_loss)
        group_loss = torch.exp(-torch.stack(group_losses)).mean()
        return group_loss     

    def forward(self,
                list_vec_groups):

        push_loss = self.loss_weight * self.group_loss(list_vec_groups)

        return push_loss

